import torch
import torch.nn as nn
import torch.nn.functional as F
from Loss.LossRegister import DefaultLossRegister as LossRegister

@LossRegister.register("CE_loss")
class CE_loss(nn.Module):
    def __init__(self):
        super().__init__()
        self.ce = nn.CrossEntropyLoss()
    def forward(self, pred, label):
        label = label.flatten()
        return self.ce(pred, label)


if __name__ == '__main__':
    x = torch.randn((3,8))
    y = torch.tensor([[1],[7],[2]])
    print(CE_loss()(x,y))